package proxyc

import (
	"context"
	"errors"
	"gitee.com/yysf-xin/dproxy/pkg/app"
	"gitee.com/yysf-xin/dproxy/pkg/db"
	"gitee.com/yysf-xin/dproxy/pkg/model"
	"gitee.com/yysf-xin/dproxy/pkg/rest"
	"gitee.com/yysf-xin/dproxy/pkg/svr"
	"gitee.com/yysf-xin/dproxy/pkg/tools"
	"gitee.com/yysf-xin/dproxy/pkg/turn"
	"github.com/fanyang01/radix"
	"github.com/go-chi/chi/v5"
	"golang.org/x/net/proxy"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	ErrServiceName = errors.New("err service name")
)

type RouteControl struct {
	DB       *db.Client
	DialRule *radix.PatternTrie
	ruleLock sync.RWMutex
	ruleMap  map[uint]*db.RouteRule
	services *app.Services

	forwardMap map[string]*db.TcpForward
	fMap       map[uint]*db.TcpForward
	fMapLock   sync.RWMutex

	reverseMap   map[string]map[string]*db.ProxyPass
	proxyMap     map[uint]*db.ProxyPass
	proxyMapLock sync.RWMutex
}

func (d *RouteControl) Dial(network, addr string) (c net.Conn, err error) {
	return d.DialContext(context.TODO(), network, addr)
}

func NewRouteControl(DB *db.Client, cfg *app.ConfigFile, services *app.Services) (r *RouteControl, err error) {

	if cfg.Db.Upgrade {
		err = db.UpgradeDB(&DB.Models, DB.Db)
		if err != nil {
			return
		}
	}

	r = &RouteControl{
		ruleMap:    make(map[uint]*db.RouteRule),
		DB:         DB,
		services:   services,
		forwardMap: make(map[string]*db.TcpForward),
		fMap:       make(map[uint]*db.TcpForward),
		reverseMap: make(map[string]map[string]*db.ProxyPass),
		proxyMap:   make(map[uint]*db.ProxyPass),
	}
	r.rebuildRule()
	tr := http.DefaultTransport.(*http.Transport)
	tr.Proxy = r.Proxy
	tr.DialContext = r.DialContext
	r.initProxyPass()
	return
}
func (d *RouteControl) FindByAddr(addr string) *db.RouteRule {
	host := addr
	if index := strings.Index(addr, ":"); index != -1 {
		host = addr[:index]
	}
	d.ruleLock.RLock()
	v, ok := d.DialRule.Lookup(host)
	d.ruleLock.RUnlock()

	if ok && v.(*db.RouteRule).Service != "" {
		return v.(*db.RouteRule)
	}
	return nil
}

func (d *RouteControl) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
	var count *int
	count, ok := ctx.Value("counter").(*int)
	if !ok {
		var ca int
		count = &ca
		ctx = context.WithValue(ctx, "counter", count)
	}
	*count++
	if *count > 10 {
		tools.Log.Println("Dial", network, addr, "times", *count, "> 10 fail back default")
		return net.DialTimeout(network, addr, 10*time.Second)
	}
	d.fMapLock.RLock()
	f := d.forwardMap[addr]
	d.fMapLock.RUnlock()

	if f != nil {
		tools.Log.Println("FORWARD", addr, "=>", f.TarIp)
		addr = f.TarIp
	}
	v := d.FindByAddr(addr)
	if v != nil {
		var sv app.Service
		if strings.HasSuffix(v.Service, turn.ServerName) {
			sv = d.services.GetService(turn.ServerName)
		} else {
			sv = d.services.GetService(v.Service)
		}
		if sv != nil {
			if c, ok := sv.(proxy.ContextDialer); ok {
				tools.Log.Println("USE", sv, "=>", addr)
				return c.DialContext(ctx, network, addr)
			}
			if c, ok := sv.(proxy.Dialer); ok {
				tools.Log.Println("USE", sv, "=>", addr)
				return c.Dial(network, addr)
			}
			tools.Log.Println("err", sv, "is not a dialer !")
		} else {
			tools.Log.Println("err find service", v.Service)
		}

	}
	tools.Log.Println("DIRECT", addr)
	return net.DialTimeout(network, addr, 10*time.Second)

}

func (d *RouteControl) Proxy(r *http.Request) (*url.URL, error) {
	var host string
	if strings.Contains(r.URL.Host, ":") {
		host, _, _ = net.SplitHostPort(r.URL.Host)
	} else {
		host = r.URL.Host
	}
	d.ruleLock.RLock()
	v, ok := d.DialRule.Lookup(host)
	d.ruleLock.RUnlock()

	if ok && v.(*db.RouteRule).Service != "" {
		sv := d.services.GetService(v.(*db.RouteRule).Service)
		if sv != nil {
			if px, ok := sv.(*svr.ProxyDialer); ok {
				return px.Proxy.ProxyURL, nil
			}
		}
	}

	return nil, nil
}

func (d *RouteControl) ChooseServerId(ctx context.Context, network, address string) uint64 {
	r := d.FindByAddr(address)
	if r == nil {
		return 0
	} else {
		return r.IntId
	}
}

func (d *RouteControl) InitForward(route chi.Router) {
	route.Delete("/{id}", app.JsonHttp(nil, d.deleteForward))
	route.Route("/{id}/", func(r chi.Router) {
		r.Post("/enable", app.JsonHttp(func() interface{} {
			return &model.Enable{}
		}, d.handleForward))
	})
	var forwards db.TcpForwardList
	forwards.ListByEnabled(d.DB.Db)
	for i := range forwards {
		d.addForward(&forwards[i])
	}
}

func (d *RouteControl) handleForward(writer http.ResponseWriter, r *http.Request, i interface{}) (rs interface{}, err error) {
	var m db.TcpForward
	err = rest.FindModelByReq(r, &m, d.DB.Db)
	if err != nil {
		return
	}
	err = i.(*model.Enable).Update(d.DB.Db.Model(&m))
	if err != nil {
		return
	}
	if i.(*model.Enable).Enabled {
		d.addForward(&m)
	} else {
		d.removeForward(&m)
	}
	return

}

func (d *RouteControl) addForward(t *db.TcpForward) {
	d.removeForward(t)

	d.fMapLock.Lock()
	d.forwardMap[t.SrcIp] = t
	d.fMap[t.Id] = t
	d.fMapLock.Unlock()

}

func (d *RouteControl) removeForward(t *db.TcpForward) {
	d.fMapLock.RLock()
	f := d.fMap[t.Id]
	d.fMapLock.RUnlock()

	if f != nil {
		d.fMapLock.Lock()
		delete(d.fMap, t.Id)
		delete(d.forwardMap, f.SrcIp)
		d.fMapLock.Unlock()
	}

}

func (d *RouteControl) deleteForward(w http.ResponseWriter, r *http.Request, i interface{}) (rs interface{}, err error) {
	var m db.TcpForward
	err = rest.FindModelByReq(r, &m, d.DB.Db)
	if err != nil {
		return
	}
	d.removeForward(&m)
	err = d.DB.Db.Delete(&m, m.Id).Error
	return
}

func (d *RouteControl) InitRule(route chi.Router) {
	route.Delete("/{id}", app.JsonHttp(nil, d.deleteRule))
	route.Route("/{id}/", func(r chi.Router) {
		r.Post("/enable", app.JsonHttp(func() interface{} {
			return &model.Enable{}
		}, d.handleRule))
	})
}

func (d *RouteControl) deleteRule(writer http.ResponseWriter, r *http.Request, i interface{}) (rs interface{}, err error) {
	var m db.RouteRule
	err = rest.FindModelByReq(r, &m, d.DB.Db)
	if err != nil {
		return
	}
	d.removeRule(&m)
	err = d.DB.Db.Delete(m).Error
	return
}

func (d *RouteControl) handleRule(w http.ResponseWriter, r *http.Request, i interface{}) (rs interface{}, err error) {
	var m db.RouteRule
	err = rest.FindModelByReq(r, &m, d.DB.Db)
	if err != nil {
		return
	}
	d.DB.Db.Model(&m).Update("enabled", i.(*model.Enable).Enabled)
	if i.(*model.Enable).Enabled {
		err = d.addRule(&m)
		rs = &m
	} else {
		d.removeRule(&m)
	}
	d.rebuildRule()
	return
}

func (d *RouteControl) removeRule(m *db.RouteRule) {
	d.ruleLock.RLock()
	delete(d.ruleMap, m.Id)
	d.ruleLock.RUnlock()
}

func (d *RouteControl) addRule(m *db.RouteRule) (err error) {
	if m.Service != "" {
		idx := strings.LastIndex(m.Service, "@")
		if idx != -1 {
			m.IntId, err = strconv.ParseUint(m.Service[:idx], 10, 64)
		} else {
			// err = ErrServiceName
		}
		if err != nil {
			return
		}
	} else {
		m.IntId = 0
	}

	d.removeRule(m)
	d.ruleLock.Lock()
	d.ruleMap[m.Id] = m
	d.DialRule.Add(m.Host, m)
	d.ruleLock.Unlock()
	return
}

func (d *RouteControl) rebuildRule() {
	r := radix.NewPatternTrie()
	var list db.RouteRuleList
	list.ListByEnabled(d.DB.Db)
	var err error
	for _, s := range list {
		err = nil
		idx := strings.LastIndex(s.Service, "@")
		if idx != -1 {
			s.IntId, err = strconv.ParseUint(s.Service[:idx], 10, 64)
		}
		if err != nil {
			tools.Log.Println(s.Host, "->", s.Id, s.Service, err)
			continue
		}
		_, has := r.Add(s.Host, s)
		if has {
			tools.Log.Println(s.Host, "->", s.Service, "repeat")
		}
	}
	d.ruleLock.Lock()
	d.DialRule = r
	d.ruleLock.Unlock()

}

func (d *RouteControl) InitProxyPassRoute(route chi.Router) {
	route.Delete("/{id}", app.JsonHttp(nil, d.deleteProxyPass))
	route.Route("/{id}/", func(r chi.Router) {
		r.Post("/enable", app.JsonHttp(func() interface{} {
			return &model.Enable{}
		}, d.handleProxyPass))
	})
}

func (d *RouteControl) initProxyPass() {
	var list []*db.ProxyPass
	d.DB.Db.Find(&list, "enabled = ?", true)
	for i := range list {
		d.addProxyPass(list[i])
	}
}

func (d *RouteControl) deleteProxyPass(_ http.ResponseWriter, r *http.Request, i interface{}) (rs interface{}, err error) {
	var m db.ProxyPass
	err = rest.FindModelByReq(r, &m, d.DB.Db)
	if err != nil {
		return
	}
	d.removeProxyPass(&m)
	return
}

func (d *RouteControl) handleProxyPass(_ http.ResponseWriter, r *http.Request, i interface{}) (rs interface{}, err error) {
	var m db.ProxyPass
	err = rest.FindModelByReq(r, &m, d.DB.Db)
	if err != nil {
		return
	}
	err = i.(*model.Enable).Update(d.DB.Db.Model(&m))
	if err != nil {
		return
	}
	if i.(*model.Enable).Enabled {
		d.proxyMapLock.Lock()
		d.addProxyPass(&m)
		d.proxyMapLock.Unlock()
		rs = &m
	} else {
		d.proxyMapLock.Lock()
		d.removeProxyPass(&m)
		d.proxyMapLock.Unlock()
	}
	return
}

func (d *RouteControl) addProxyPass(pass *db.ProxyPass) {
	d.removeProxyPass(pass)
	p := d.reverseMap[pass.Domain]
	if p == nil {
		p = make(map[string]*db.ProxyPass)
		d.reverseMap[pass.Domain] = p
	}
	p[pass.Path] = pass
	d.proxyMap[pass.Id] = pass
}
func (d *RouteControl) removeProxyPass(pass *db.ProxyPass) {
	p := d.proxyMap[pass.Id]
	if p == nil {
		return
	}
	m := d.reverseMap[p.Domain]
	if m == nil {
		return
	}
	px := m[p.Path]
	if px != nil {
		delete(m, p.Path)
	}

	if len(m) == 0 {
		delete(d.reverseMap, p.Domain)
	}

}
func (d *RouteControl) HasReverse(host string) bool {
	return d.reverseMap[host] != nil
}
func (d *RouteControl) Director(r *http.Request) {
	d.proxyMapLock.RLock()
	defer d.proxyMapLock.RUnlock()
	if r.TLS != nil && r.URL.Host == "" {
		r.URL.Host = r.Host
		r.URL.Scheme = "https"
	}
	m := d.reverseMap[r.URL.Host]
	if m == nil {
		return
	}

	for uri, proxyCfg := range m {
		if !strings.HasPrefix(r.URL.Path, uri) {
			continue
		}
		if proxyCfg.TargetURL == nil {
			proxyCfg.TargetURL, _ = url.Parse(proxyCfg.Target)
		}
		if proxyCfg.TargetURL == nil {
			tools.Log.Printf("Invalid URL %s\n", proxyCfg.Target)
		} else {
			target := proxyCfg.TargetURL
			targetQuery := target.RawQuery
			r.URL.Scheme = target.Scheme
			r.URL.Host = target.Host
			if proxyCfg.ChangeOrigin {
				r.Host = target.Host
			}
			if strings.HasSuffix(proxyCfg.Target, "/") {
				r.URL.Path = tools.SingleJoiningSlash(target.Path, r.URL.Path[len(uri):])
			} else {
				r.URL.Path = tools.SingleJoiningSlash(target.Path, r.URL.Path)
			}
			host, _, _ := net.SplitHostPort(target.Host)
			ip := net.ParseIP(host)
			if ip == nil {
				// 目标不是ip再重写host
				r.Host = target.Host
			}
			if targetQuery == "" || r.URL.RawQuery == "" {
				r.URL.RawQuery = targetQuery + r.URL.RawQuery
			} else {
				r.URL.RawQuery = targetQuery + "&" + r.URL.RawQuery
			}
			//r.Header.Set("Accept-Encoding", "identity")
			tools.Log.Println("PROXY_PASS", proxyCfg.Domain, proxyCfg.Target)
			return
		}
	}
}
